"""
实验记录器模块
该模块提供了一个Recorder类，用于管理实验结果的保存、代码备份和日志记录
"""

# 导入操作系统接口模块，提供文件和目录操作功能
import os
# 导入高级文件操作模块，提供文件复制、移动等操作
import shutil


# 定义Recorder类，继承自object基类
class Recorder(object):
    """
    实验记录器类
    用于管理实验过程中的文件保存、代码备份、日志记录等功能
    """
    
    def __init__(self, snapshot_pref, ignore_folder):
        """
        初始化记录器
        
        参数:
        snapshot_pref: 快照保存路径前缀，实验相关文件将保存在此路径下
        ignore_folder: 需要忽略的文件夹名称，在代码备份时会跳过包含此名称的目录
        """
        
        # 检查快照保存路径是否存在，如果不存在则创建
        if not os.path.isdir(snapshot_pref):
            os.mkdir(snapshot_pref)
        
        # 设置保存路径
        self.save_path = snapshot_pref
        # 日志文件路径
        self.log_file = self.save_path + "log.txt"
        # README文件路径
        self.readme = self.save_path + "README.md"
        # 参数配置文件路径
        self.opt_file = self.save_path + "opt.log"
        # 代码备份路径
        self.code_path = os.path.join(self.save_path, "code/")
        
        # 以下是被注释掉的权重相关路径设置
        # self.weight_folder = self.save_path + "weight/"  # 权重保存文件夹
        # self.weight_fig_folder = self.save_path + "weight_fig/"  # 权重可视化图片文件夹
        
        # 以下是被注释掉的日志文件删除代码
        # if os.path.isfile(self.log_file):
        #     os.remove(self.log_file)
        
        # 如果README文件已存在，则删除（重新开始记录）
        if os.path.isfile(self.readme):
            os.remove(self.readme)
        
        # 如果代码备份路径不存在，则创建
        if not os.path.isdir(self.code_path):
            os.mkdir(self.code_path)
        
        # 复制源代码到备份路径，忽略指定文件夹
        self.copy_code(dst=self.code_path, ignore_folder=ignore_folder)
        
        # 以下是被注释掉的权重文件夹管理代码
        """if os.path.isdir(self.weight_folder):
            shutil.rmtree(self.weight_folder, ignore_errors=True)
        os.mkdir(self.weight_folder)
        if os.path.isdir(self.weight_fig_folder):
            shutil.rmtree(self.weight_fig_folder, ignore_errors=True)
        os.mkdir(self.weight_fig_folder)"""

        # 打印保存路径信息
        print ("\n======> Result will be saved at: ", self.save_path)

    def copy_code(self, src="./", dst="./code/", ignore_folder='Exps'):
        """
        复制源代码文件到指定目录
        
        参数:
        src: 源代码根目录路径，默认为当前目录"./"
        dst: 目标备份目录路径，默认为"./code/"
        ignore_folder: 需要忽略的文件夹名称，会跳过包含此名称的目录
        
        功能: 遍历源代码目录，将所有符合条件的文件复制到备份目录
        """
        
        # 导入UUID模块，用于生成唯一标识符
        import uuid
        
        # 如果目标目录已存在，则生成带有UUID的新目录名以避免冲突
        if os.path.isdir(dst):
            # 以下是被注释掉的另一种目录命名方式
            # dst = "/".join(dst.split("/")[:-1])+"_"+str(uuid.uuid4())+"/"
            
            # 生成带有UUID的新目录路径
            dst = "/".join(dst.split("/")) + "code_" + str(uuid.uuid4()) + "/"
        
        # 初始化文件绝对路径列表
        file_abs_list = []
        # 获取源代码目录的绝对路径
        src_abs = os.path.abspath(src)
        
        # 遍历源代码目录树
        for root, dirs, files in os.walk(src_abs):
            # 跳过包含忽略文件夹的目录
            if ignore_folder not in root:
                # 遍历当前目录下的所有文件
                for name in files:
                    # 将文件的绝对路径添加到列表中
                    file_abs_list.append(root + "/" + name)
        
        # 遍历所有找到的文件
        for file_abs in file_abs_list:
            # 分割文件名和扩展名
            file_split = file_abs.split("/")[-1].split('.')
            
            # 以下是被注释掉的仅复制Python文件的代码
            # if len(file_split) >= 2 and file_split[1] == "py":
            
            # 检查文件大小是否小于10MB且不是.pyc文件
            if os.path.getsize(file_abs)/1024/1024 < 10 and not file_split[-1] == "pyc":
                # 源文件路径
                src_file = file_abs
                # 目标文件路径，保持相对目录结构
                dst_file = dst + file_abs.replace(src_abs, "")
                
                # 如果目标文件的目录不存在，则创建
                if not os.path.exists(os.path.dirname(dst_file)):
                    os.makedirs(os.path.dirname(dst_file))
                
                # 复制文件（这里有两行相同的复制代码，可能是重复的）
                shutil.copyfile(src=src_file, dst=dst_file)
                
                # 尝试复制文件，如果出错则打印错误信息
                try:
                    shutil.copyfile(src=src_file, dst=dst_file)
                except:
                    print("copy file error")

    def writeopt(self, opt):
        """
        写入实验参数到文件
        
        参数:
        opt: 包含实验参数的对象，通常是一个argparse.Namespace对象
        
        功能: 将实验参数以键值对形式写入文件
        """
        
        # 以写入模式打开参数文件
        with open(self.opt_file, "w") as f:
            # 遍历参数对象的所有属性
            for k, v in opt.__dict__.items():
                # 将参数名和值写入文件
                f.write(str(k)+": "+str(v)+"\n")

    def writelog(self, input_data):
        """
        写入日志信息
        
        参数:
        input_data: 要写入日志的数据，可以是任意类型，会被转换为字符串
        
        功能: 以追加模式将数据写入日志文件
        """
        
        # 以追加模式打开日志文件
        txt_file = open(self.log_file, 'a+')
        # 写入数据并换行
        txt_file.write(str(input_data) + "\n")
        # 关闭文件
        txt_file.close()

    def writereadme(self, input_data):
        """
        写入README信息
        
        参数:
        input_data: 要写入README文件的数据
        
        功能: 以追加模式将数据写入README文件
        """
        
        # 以追加模式打开README文件
        txt_file = open(self.readme, 'a+')
        # 写入数据并换行
        txt_file.write(str(input_data) + "\n")
        # 关闭文件
        txt_file.close()


    def gennetwork(self, var):
        """
        生成网络结构图（方法未完整实现）
        
        参数:
        var: 用于生成网络图的变量或数据
        
        注意: 此方法调用了self.graph.draw，但graph属性在类中未定义
        """
        
        # 调用graph对象的draw方法绘制网络图
        self.graph.draw(var=var)

    def savenetwork(self):
        """
        保存网络结构图（方法未完整实现）
        
        功能: 将网络图保存为SVG格式文件
        
        注意: 此方法调用了self.graph.save，但graph属性在类中未定义
        """
        
        # 调用graph对象的save方法保存网络图
        self.graph.save(file_name=self.save_path+"network.svg")

    # 以下是被注释掉的权重记录和可视化方法
    """def writeweights(self, input_data, block_id, layer_id, epoch_id):
        # 写入权重数据到文件
        # 参数:
        # input_data: 权重数据
        # block_id: 网络块ID
        # layer_id: 网络层ID  
        # epoch_id: 训练轮次ID
        
        txt_path = self.weight_folder + "conv_weight_" + str(epoch_id) + ".log"
        txt_file = open(txt_path, 'a+')
        write_str = "%d\t%d\t%d\t" % (epoch_id, block_id, layer_id)
        for x in input_data:
            write_str += str(x) + "\t"
        txt_file.write(write_str+"\n")

    def drawhist(self):
        # 绘制权重直方图
        # 功能: 根据权重数据绘制直方图
        
        drawer = DrawHistogram(txt_folder=self.weight_folder, fig_folder=self.weight_fig_folder)
        drawer.draw()"""